from net_train import Net
import torch
import image_split
from torchvision import transforms
import net_data as data
import os

def preprocess(img_path):
    image_split.setSaveFileMode(False)
    images = image_split.split_image(img_path)
    nor_img = []
    for img in images:
        norl = transforms.Normalize(img.mean(), img.std())
        img_t = torch.from_numpy(img.reshape(1, 40, 40)).type(dtype=torch.float32)
        nor_img.append(norl(img_t))
    return nor_img

#启动初始化模型
net = Net()
sd = torch.load("./model/net_cpu.pt")
net.load_state_dict(sd)


def usage(img_path):
    imgs = preprocess(img_path)
    labels = []
    for img in imgs:
        weight_lab = net(img.reshape(1, 1, 40, 40))
        max_idx = torch.argmax(weight_lab).item()
        labels.append(chr(data.labels[max_idx]))
    print(labels)
    # os.rename(img_path, "{}/{}.jpg".format(os.path.dirname(img_path), "".join(labels)))
    return labels


if (__name__ == "__main__"):
    path = ".\\usage_src_img\\"
    l = os.listdir(path)
    for f in l:
        usage("{}{}".format(path,f))